import backtrader as bt
import pandas as pd
from datetime import datetime

# ==================== 配置表格 ====================
# 在这里修改参数，无需深入代码内部

CONFIG = {
    'data_file': 'maotai_data_600519_20230101_20250901.csv',  # 数据文件名
    'start_date': datetime(2023, 1, 1),                    # 回测开始时间
    'end_date': datetime(2025, 9, 1),                    # 回测结束时间
    'initial_cash': 100000,                               # 初始资金
    'rsi_period': 14,                                     # RSI周期
    'rsi_oversold': 30,                                   # RSI超卖线
    'rsi_overbought': 50,                                 # RSI超买线
}
# =================================================

class RSIStrategy(bt.Strategy):
    params = (
        ('rsi_period', CONFIG['rsi_period']),
        ('rsi_oversold', CONFIG['rsi_oversold']),
        ('rsi_overbought', CONFIG['rsi_overbought']),
    )
    
    def __init__(self):
        # 初始化RSI指标
        self.rsi = bt.indicators.RSI(
            self.data.close, 
            period=self.params.rsi_period
        )
        self.order = None
        self.buy_price = None
        self.buy_date = None
        
    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # 订单已提交/接受 - 无需操作
            return
            
        if order.status in [order.Completed]:
            if order.isbuy():
                self.buy_price = order.executed.price
                self.buy_date = self.data.datetime.date(0)
                print(f'{self.buy_date} - 买入执行: 价格={self.buy_price:.2f}, 成本={order.executed.value:.2f}, 佣金={order.executed.comm:.2f}')
            else:  # 卖出
                profit = order.executed.pnl
                profit_percent = (order.executed.price - self.buy_price) / self.buy_price * 100
                print(f'{self.data.datetime.date(0)} - 卖出执行: 价格={order.executed.price:.2f}, '
                      f'盈利={profit:.2f}元, 收益率={profit_percent:.2f}%')
                
                # 重置买入信息
                self.buy_price = None
                self.buy_date = None
                
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            print('订单取消/保证金不足/被拒绝')
            
        # 重置订单
        self.order = None
        
    def next(self):
        current_date = self.data.datetime.date(0)
        current_close = self.data.close[0]
        current_rsi = self.rsi[0]
        
        # 如果已有持仓
        if self.position:
            # RSI高于超买线，卖出
            if current_rsi > self.params.rsi_overbought:
                print(f'{current_date} - 卖出信号: RSI={current_rsi:.2f}, 当前价格={current_close:.2f}')
                self.sell()
                
        # 如果没有持仓
        else:
            # RSI低于超卖线，买入
            if current_rsi < self.params.rsi_oversold:
                print(f'{current_date} - 买入信号: RSI={current_rsi:.2f}, 当前价格={current_close:.2f}')
                self.buy()

class BacktestEngine:
    def __init__(self, data_path):
        self.cerebro = bt.Cerebro()
        self.data_path = data_path
        
    def load_data(self):
        # 加载数据到Pandas DataFrame
        df = pd.read_csv(self.data_path, parse_dates=['date'], index_col='date')
        
        # 创建Backtrader数据源
        data = bt.feeds.PandasData(
            dataname=df,
            open='open',
            high='high',
            low='low',
            close='close',
            volume='volume',
            fromdate=CONFIG['start_date'],
            todate=CONFIG['end_date']
        )
        
        self.cerebro.adddata(data)
        return data
        
    def add_strategy(self, strategy, **kwargs):
        self.cerebro.addstrategy(strategy, **kwargs)
        
    def set_cash(self, cash):
        self.cerebro.broker.set_cash(cash)
        
    def add_analyzer(self):
        self.cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
        self.cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
        self.cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
        
    def run_backtest(self):
        # 设置初始资金
        self.set_cash(CONFIG['initial_cash'])
        
        # 添加分析器
        self.add_analyzer()
        
        # 运行回测
        results = self.cerebro.run()
        
        # 打印结果
        strat = results[0]
        print('最终资金: %.2f' % self.cerebro.broker.getvalue())
        print('夏普比率:', strat.analyzers.sharpe.get_analysis()['sharperatio'])
        print('最大回撤: %.2f' % strat.analyzers.drawdown.get_analysis()['max']['drawdown'])
        print('年化回报: %.2f' % (strat.analyzers.returns.get_analysis()['rnorm100']))
        
        return results
        
    def plot_result(self):
        self.cerebro.plot(style='candlestick')

# 主函数
if __name__ == '__main__':
    # 创建回测引擎
    engine = BacktestEngine(CONFIG['data_file'])
    
    # 加载数据
    engine.load_data()
    
    # 添加策略
    engine.add_strategy(RSIStrategy)
    
    # 运行回测
    results = engine.run_backtest()
    
    # 绘制图表
    engine.plot_result()